from datetime import datetime
from pathlib import Path
from typing import List, Generator

import fitz


def get_timestamp() -> str:
    now = datetime.now()
    date_string = now.strftime("%Y%m%d")
    time_string = now.strftime("%H%M%S")
    return date_string + "." + time_string

def find_files_with_extensions(directory: str, extensions: dict) -> Generator[Path, None, None]:
    path = Path(directory).expanduser()
    for file_path in path.glob("*"):
        if file_path.suffix.lower() in extensions:
            yield file_path  # Use Generator return matched file path

def merge_pdf(doc_list: List[str], out_pdf: str) -> None:
    with fitz.open() as target:
        for doc in doc_list:
            with fitz.open(doc) as insert_doc:
                target.insert_pdf(insert_doc)
        target.save(out_pdf)

def delete_pages(source_path: str, out_pdf: str, f_page: int, t_page: int) -> None:
    with fitz.open(source_path) as doc:
        doc.delete_pages(from_page=f_page-1, to_page=t_page-1)
        doc.save(out_pdf)

def insert_page_with_blank(source_pdf: str, output_pdf: str, page_number: int) -> None:
    with fitz.open(source_pdf) as doc:
        doc.new_page(page_number-1, width = 595, height = 842)
        doc.save(output_pdf)

def insert_pages(main_pdf, slave_pdf, output_pdf, pos):
    pos -= 1
    doc_1 = fitz.open(main_pdf)
    doc_2 = fitz.open(slave_pdf)
    result = fitz.open()
    
    page_1 = [page for page in range(pos)]
    page_3 = [page for page in range(pos, doc_1.page_count)]
    doc_1.select(page_1)
    doc_3 = fitz.open(main_pdf)
    doc_3.select(page_3)

    result.insert_pdf(doc_1)
    result.insert_pdf(doc_2)
    result.insert_pdf(doc_3)
    result.save(output_pdf)

    doc_1.close()
    doc_2.close()
    doc_3.close()
    result.close()

def extract_pages(source_pdf: str, output_pdf: str, pages: List[int]) -> None:
    '''
    :params: 
        pages = [1, 3, 12, 11]
        pages = [4]
        pages = [x for x in range(4, 10)]
    '''
    pages = [page - 1 for page in pages]
    with fitz.open(source_pdf) as doc:
        doc.select(pages)
        doc.save(output_pdf)

def pdf_to_img(source_pdf: str, output_img: str, type: str, dpi: int, pages: List[int]) -> None:
    '''
    :params: 
        type: jpeg/png/...
        pages = [1, 3, 12, 11] or [4] or [1]
    '''
    doc = fitz.open(source_pdf)
    for page_num in pages:
        if page_num <= len(doc):
            page = doc[page_num-1]
            pix = page.get_pixmap(matrix=fitz.Matrix(300/72), dpi=dpi)
            output_file = f"{output_img}_{page_num}.{type}"
            pix.save(output_file)
    doc.close()

def images_to_pdf(image_paths: List[str], output_path: str) -> None:
    pdf = fitz.open()
    for image_path in image_paths:
        img = fitz.open(image_path)
        pdf_bytes = img.convert_to_pdf()
        img.close()
        img_pdf = fitz.open("pdf", pdf_bytes)
        pdf.insert_pdf(img_pdf)
        img_pdf.close()
    pdf.save(output_path)
    pdf.close()

def sort_toc(toc: List[List]) -> List[List]:
    return sorted(toc, key=lambda x: (x[2], -x[3] if len(x) > 3 else float('-inf')))

def bookmarks_add(pdf_source: str, pdf_target: str, new_items: list) -> None:
    with fitz.open(pdf_source) as doc:
        toc_old = doc.get_toc()
        
        for new_item in new_items:
            toc_old.append(new_item)
        
        toc_new = sort_toc(toc_old)

        doc.set_toc(toc_new)
        doc.save(pdf_target)

def delete_bookmarks_based_on_level(toc: List[List], title: str) -> bool:
    for i in range(len(toc)):
        if toc[i][1] == title:
            level = toc[i][0]
            j = i + 1
            while j < len(toc) and toc[j][0] > level:
                del toc[j]
            del toc[i]
            return True
    return False

def bookmarks_delete_by_title(pdf_source: str, pdf_target: str, titles: list) -> None:
    with fitz.open(pdf_source) as doc:
        toc_old = doc.get_toc()
        
        for title in titles:
            if not delete_bookmarks_based_on_level(toc_old, title):
                print(f"Bookmarks did't find the title: '{title}'.")
        
        doc.set_toc(toc_old)
        doc.save(pdf_target)

def main_delete_pages():
    '''Delete Pages'''
    print('\n *** Delete PDF Pages ***')
    source_path = 'path/PDF.pdf'
    out_pdf = f'path/result.{get_timestamp()}.pdf'

    # if you want delete single page, let first_page == end_page
    first_page = 3
    end_page = 4
    delete_pages(source_path, out_pdf, f_page=first_page, t_page=end_page)
    print('*** Delete PDF pages SUCCESS ***\n')

def main_insert_pages():
    '''Insert Pages'''
    # 这里有个问题还没解决：目录丢失问题
    print('\n *** Insert PDF Pages ***')
    source_path = 'path/PDF_main.pdf'
    slave_path = 'path/PDF_slave.pdf'
    insert_pos = 5
    out_pdf = f'path/result.{get_timestamp()}.pdf'
    insert_pages(source_path, slave_path, out_pdf, insert_pos)
    print('*** Insert PDF pages SUCCESS ***\n')

def main_insert_blank_page():
    '''Insert Pages with Blank'''
    print('\n *** Insert PDF with Blank page ***')
    source_path = 'path/1.pdf'
    out_pdf = f'path/result.{get_timestamp()}.pdf'
    page_number = 3
    insert_page_with_blank(source_path, out_pdf, page_number)
    print('*** Insert PDF pages SUCCESS ***\n')

def main_extract_pages():
    '''Extract Pages'''
    print('\n *** Extract PDF pages form a single PDF ***')
    source_path = './math.pdf'
    out_pdf = f'./math.{get_timestamp()}.pdf'
    
    # pages = [1, 2, 3, 4, 5, 6]
    # pages = [2]
    pages = [x for x in range(188, 407)]
    
    extract_pages(source_path, out_pdf, pages)
    print('*** Extract PDF pages SUCCESS ***\n')

def main_image2pdf():
    '''Convert Images to PDF'''
    print('\n *** Convert Images to PDF ***')
    directory = 'path'
    extensions = {'.jpg'}
    
    doc_list = sorted(find_files_with_extensions(directory, extensions), key=lambda p: p.name)
    out_pdf = f'./result.{get_timestamp()}.pdf'
    images_to_pdf(doc_list, out_pdf)
    print('*** Convert Images to PDF SUCCESS ***\n')

def main_pdf2image():
    '''Convert PDF to Image(jpeg, png, ...)'''
    print('\n *** Convert PDF to Images ***')
    source_path = 'path/PDF.pdf'
    out_img = f'path/result.{get_timestamp()}'
    # pages = [1, 3, 4]
    pages = [x for x in range(1, 9)]
    filetype = 'png'
    pdf_to_img(source_path, out_img, filetype, dpi=300, pages=pages)
    print('\n *** Convert PDF to Images SUCCESS ***')

def main_merge_pdf():
    '''Merge PDF list'''
    # TODO: lose bookmarks
    print('\n *** Merge multi PDFs ***')
    directory = "/Users/kxz/Downloads"
    # give the wanted file extensions, eg. {".pdf", ".py", '.txt'}
    extensions = {".pdf"}
    out_pdf = f'path/result.{get_timestamp()}.pdf'
    doc_list = sorted(find_files_with_extensions(directory, extensions), key=lambda p: p.name)
    merge_pdf(doc_list=doc_list, out_pdf=out_pdf)
    print('\n *** Merge multi PDFs SUCCESS ***')

def main_bookmarks_add():
    '''Add Bookmarks'''
    pdf_source = './temp.pdf'
    pdf_target = './temp_new.pdf'

    new_items = [
        [1, 'title 232422', 11, 130],
        [2, 'title new2', 5, 140],
        [2, 'title new3', 6, 150]
    ]
    
    bookmarks_add(pdf_source, pdf_target, new_items)
    print('\n *** Add New Bookmarks SUCCESS ***')

def main_bookmarks_delete_by_title():
    '''Delete Bookmarks by Title'''
    pdf_source = 'path/PDF.pdf'
    pdf_target = 'path/PDF_temp.pdf'
    
    titles = [
        '1 Introduction ',
        'Appendix A Facts from Linear Algebra and Matrix Analysis ',
        'Appendix B Convex Sets and Functions ',
        'Appendix C Optimization Problems and Optimality Conditions ',
        'Appendix D Methods for Optimization ',
        'Appendix E Facts from High-Dimensional Statistics '
    ]
    
    bookmarks_delete_by_title(pdf_source, pdf_target, titles)
    print('Delete Bookmarks by Title SUCCESS...')

if __name__ == '__main__':
    '''PDF pages'''
    # main_delete_pages()
    # main_insert_pages()
    # main_insert_blank_page()
    # main_extract_pages()
    # main_merge_pdf()

    '''PDF images'''
    # main_image2pdf()
    # main_pdf2image()

    '''PDF bookmarks'''
    # main_bookmarks_add()
    # main_bookmarks_delete_by_title()
    bookmarks_get()
